{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chinese-abortion",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "import numpy as np\n",
    "import itertools\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torch.optim as optim\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "def tsp_opt(points):\n",
    "    \"\"\"\n",
    "    Dynamic programing solution for TSP - O(2^n*n^2)\n",
    "    https://gist.github.com/mlalevic/6222750\n",
    "\n",
    "    :param points: List of (x, y) points\n",
    "    :return: Optimal solution\n",
    "    \"\"\"\n",
    "\n",
    "    def length(x_coord, y_coord):\n",
    "        return np.linalg.norm(np.asarray(x_coord) - np.asarray(y_coord))\n",
    "\n",
    "    # Calculate all lengths\n",
    "    all_distances = [[length(x, y) for y in points] for x in points]\n",
    "    # Initial value - just distance from 0 to every other point + keep the track of edges\n",
    "    A = {(frozenset([0, idx+1]), idx+1): (dist, [0, idx+1]) for idx, dist in enumerate(all_distances[0][1:])}\n",
    "    cnt = len(points)\n",
    "    for m in range(2, cnt):\n",
    "        B = {}\n",
    "        for S in [frozenset(C) | {0} for C in itertools.combinations(range(1, cnt), m)]:\n",
    "            for j in S - {0}:\n",
    "                # This will use 0th index of tuple for ordering, the same as if key=itemgetter(0) used\n",
    "                B[(S, j)] = min([(A[(S-{j}, k)][0] + all_distances[k][j], A[(S-{j}, k)][1] + [j])\n",
    "                                 for k in S if k != 0 and k != j])\n",
    "        A = B\n",
    "\n",
    "    res = min([(A[d][0] + all_distances[0][d[1]], A[d][1]) for d in iter(A)])\n",
    "    return np.asarray(res[1])\n",
    "\n",
    "\n",
    "class TSPDataset(Dataset):\n",
    "    \"\"\"\n",
    "    Random TSP dataset\n",
    "    \"\"\"\n",
    "    def __init__(self, data_size, seq_len, solver=tsp_opt, solve=True):\n",
    "        self.data_size = data_size\n",
    "        self.seq_len = seq_len \n",
    "        self.solve = solve\n",
    "        self.solver = solver\n",
    "        self.data = self._generate_data()\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data_size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        tensor = torch.from_numpy(self.data['Points_List'][idx]).float()\n",
    "        solution = torch.from_numpy(self.data['Solutions'][idx]).long() if self.solve else None\n",
    "\n",
    "        sample = {'Points': tensor, 'Solution': solution}\n",
    "\n",
    "        return sample\n",
    "\n",
    "    def _generate_data(self):\n",
    "        \"\"\"\n",
    "        :return: Set of points_list ans their One-Hot vector solutions\n",
    "        \"\"\"\n",
    "        points_list = []\n",
    "        solutions = []\n",
    "        data_iter = tqdm(range(self.data_size), unit='data')\n",
    "        for i, _ in enumerate(data_iter):\n",
    "            data_iter.set_description('Data points %i/%i' % (i+1, self.data_size))\n",
    "            points_list.append(np.random.random((self.seq_len, 2)))\n",
    "        solutions_iter = tqdm(points_list, unit='solve')\n",
    "        if self.solve:\n",
    "            #\n",
    "            for i, points in enumerate(solutions_iter):\n",
    "                solutions_iter.set_description('Solved %i/%i' % (i+1, len(points_list)))\n",
    "                solutions.append(self.solver(points))\n",
    "        else:\n",
    "            solutions = None\n",
    "\n",
    "        return {'Points_List': points_list, 'Solutions': solutions}\n",
    "\n",
    "    def _to1hotvec(self, points):\n",
    "        \"\"\"\n",
    "        :param points: List of integers representing the points indexes\n",
    "        :return: Matrix of One-Hot vectors\n",
    "        \"\"\"\n",
    "        vec = np.zeros((len(points), self.seq_len))\n",
    "        for i, v in enumerate(vec):\n",
    "            v[points[i]] = 1\n",
    "\n",
    "        return vec\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    \"\"\"\n",
    "    Encoder class for Pointer-Net\n",
    "    \"\"\"\n",
    "    def __init__(self, embedding_dim,\n",
    "                 hidden_dim,\n",
    "                 n_layers,\n",
    "                 dropout,\n",
    "                 bidir):\n",
    "        \"\"\"\n",
    "        Initiate Encoder\n",
    "\n",
    "        :param Tensor embedding_dim: Number of embbeding channels\n",
    "        :param int hidden_dim: Number of hidden units for the LSTM\n",
    "        :param int n_layers: Number of layers for LSTMs\n",
    "        :param float dropout: Float between 0-1\n",
    "        :param bool bidir: Bidirectional\n",
    "        \"\"\"\n",
    "        \n",
    "        super(Encoder, self).__init__()\n",
    "        self.hidden_dim = hidden_dim//2 if bidir else hidden_dim\n",
    "        self.n_layers = n_layers*2 if bidir else n_layers\n",
    "        self.bidir = bidir\n",
    "        self.lstm = nn.LSTM(embedding_dim,\n",
    "                            self.hidden_dim,\n",
    "                            n_layers,\n",
    "                            dropout=dropout,\n",
    "                            bidirectional=bidir)\n",
    "\n",
    "        # Used for propagating .cuda() command\n",
    "        self.h0 = Parameter(torch.zeros(1), requires_grad=False)\n",
    "        self.c0 = Parameter(torch.zeros(1), requires_grad=False)\n",
    "\n",
    "    def forward(self, embedded_inputs,\n",
    "                hidden):\n",
    "        \"\"\"\n",
    "        Encoder - Forward-pass\n",
    "\n",
    "        :param Tensor embedded_inputs: Embedded inputs of Pointer-Net\n",
    "        :param Tensor hidden: Initiated hidden units for the LSTMs (h, c)\n",
    "        :return: LSTMs outputs and hidden units (h, c)\n",
    "        \"\"\"\n",
    "        \n",
    "        embedded_inputs = embedded_inputs.permute(1, 0, 2)\n",
    "\n",
    "        outputs, hidden = self.lstm(embedded_inputs, hidden)\n",
    "\n",
    "        return outputs.permute(1, 0, 2), hidden\n",
    "\n",
    "    def init_hidden(self, embedded_inputs):\n",
    "        \"\"\"\n",
    "        Initiate hidden units\n",
    "\n",
    "        :param Tensor embedded_inputs: The embedded input of Pointer-NEt\n",
    "        :return: Initiated hidden units for the LSTMs (h, c)\n",
    "        \"\"\"\n",
    "\n",
    "        batch_size = embedded_inputs.size(0)\n",
    "\n",
    "        # Reshaping (Expanding)\n",
    "        h0 = self.h0.unsqueeze(0).unsqueeze(0).repeat(self.n_layers,\n",
    "                                                      batch_size,\n",
    "                                                      self.hidden_dim)\n",
    "        c0 = self.h0.unsqueeze(0).unsqueeze(0).repeat(self.n_layers,\n",
    "                                                      batch_size,\n",
    "                                                      self.hidden_dim)\n",
    "\n",
    "        return h0, c0\n",
    "\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    \"\"\"\n",
    "    Attention model for Pointer-Net\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_dim,\n",
    "                 hidden_dim):\n",
    "        \"\"\"\n",
    "        Initiate Attention\n",
    "\n",
    "        :param int input_dim: Input's diamention\n",
    "        :param int hidden_dim: Number of hidden units in the attention\n",
    "        \"\"\"\n",
    "\n",
    "        super(Attention, self).__init__()\n",
    "\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.input_linear = nn.Linear(input_dim, hidden_dim)\n",
    "        self.context_linear = nn.Conv1d(input_dim, hidden_dim, 1, 1)\n",
    "        self.V = Parameter(torch.FloatTensor(hidden_dim), requires_grad=True)\n",
    "        self._inf = Parameter(torch.FloatTensor([float('-inf')]), requires_grad=False)\n",
    "        self.tanh = nn.Tanh()\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "\n",
    "        # Initialize vector V\n",
    "        nn.init.uniform_(self.V, -1, 1)\n",
    "\n",
    "    def forward(self, input,\n",
    "                context,\n",
    "                mask):\n",
    "        \"\"\"\n",
    "        Attention - Forward-pass\n",
    "\n",
    "        :param Tensor input: Hidden state h\n",
    "        :param Tensor context: Attention context\n",
    "        :param ByteTensor mask: Selection mask\n",
    "        :return: tuple of - (Attentioned hidden state, Alphas)\n",
    "        \"\"\"\n",
    "        \n",
    "        # (batch, hidden_dim, seq_len)\n",
    "        inp = self.input_linear(input).unsqueeze(2).expand(-1, -1, context.size(1))\n",
    "\n",
    "        # (batch, hidden_dim, seq_len)\n",
    "        context = context.permute(0, 2, 1)\n",
    "        ctx = self.context_linear(context)\n",
    "\n",
    "        # (batch, 1, hidden_dim)\n",
    "        V = self.V.unsqueeze(0).expand(context.size(0), -1).unsqueeze(1)\n",
    "\n",
    "        # (batch, seq_len)\n",
    "        att = torch.bmm(V, self.tanh(inp + ctx)).squeeze(1)\n",
    "        if len(att[mask]) > 0:\n",
    "            att[mask] = self.inf[mask]\n",
    "        alpha = self.softmax(att)\n",
    "\n",
    "        hidden_state = torch.bmm(ctx, alpha.unsqueeze(2)).squeeze(2)\n",
    "\n",
    "        return hidden_state, alpha\n",
    "\n",
    "    def init_inf(self, mask_size):\n",
    "        self.inf = self._inf.unsqueeze(1).expand(*mask_size)\n",
    "\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    \"\"\"\n",
    "    Decoder model for Pointer-Net\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embedding_dim,\n",
    "                 hidden_dim):\n",
    "        \"\"\"\n",
    "        Initiate Decoder\n",
    "\n",
    "        :param int embedding_dim: Number of embeddings in Pointer-Net\n",
    "        :param int hidden_dim: Number of hidden units for the decoder's RNN\n",
    "        \"\"\"\n",
    "        \n",
    "        super(Decoder, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.input_to_hidden = nn.Linear(embedding_dim, 4 * hidden_dim)\n",
    "        self.hidden_to_hidden = nn.Linear(hidden_dim, 4 * hidden_dim)\n",
    "        self.hidden_out = nn.Linear(hidden_dim * 2, hidden_dim)\n",
    "        \n",
    "        self.att = Attention(hidden_dim, hidden_dim)\n",
    "\n",
    "        # Used for propagating .cuda() command\n",
    "        self.mask = Parameter(torch.ones(1), requires_grad=False)\n",
    "        self.runner = Parameter(torch.zeros(1), requires_grad=False)\n",
    "\n",
    "    def forward(self, embedded_inputs,\n",
    "                decoder_input,\n",
    "                hidden,\n",
    "                context):\n",
    "        \"\"\"\n",
    "        Decoder - Forward-pass\n",
    "\n",
    "        :param Tensor embedded_inputs: Embedded inputs of Pointer-Net\n",
    "        :param Tensor decoder_input: First decoder's input\n",
    "        :param Tensor hidden: First decoder's hidden states\n",
    "        :param Tensor context: Encoder's outputs\n",
    "        :return: (Output probabilities, Pointers indices), last hidden state\n",
    "        \"\"\"\n",
    "        \n",
    "        batch_size = embedded_inputs.size(0)\n",
    "        input_length = embedded_inputs.size(1)\n",
    "\n",
    "        # (batch, seq_len)\n",
    "        mask = self.mask.repeat(input_length).unsqueeze(0).repeat(batch_size, 1)\n",
    "        self.att.init_inf(mask.size())\n",
    "\n",
    "        # Generating arang(input_length), broadcasted across batch_size\n",
    "        runner = self.runner.repeat(input_length)\n",
    "        for i in range(input_length):\n",
    "            runner.data[i] = i\n",
    "        runner = runner.unsqueeze(0).expand(batch_size, -1).long()\n",
    "\n",
    "        outputs = []\n",
    "        pointers = []\n",
    "\n",
    "        def step(x, hidden):\n",
    "            \"\"\"\n",
    "            Recurrence step function\n",
    "\n",
    "            :param Tensor x: Input at time t\n",
    "            :param tuple(Tensor, Tensor) hidden: Hidden states at time t-1\n",
    "            :return: Hidden states at time t (h, c), Attention probabilities (Alpha)\n",
    "            \"\"\"\n",
    "\n",
    "            # Regular LSTM\n",
    "            h, c = hidden\n",
    "\n",
    "            gates = self.input_to_hidden(x) + self.hidden_to_hidden(h)\n",
    "            input, forget, cell, out = gates.chunk(4, 1)\n",
    "\n",
    "            input = torch.sigmoid(input)\n",
    "            forget = torch.sigmoid(forget)\n",
    "            cell = torch.tanh(cell)\n",
    "            out = torch.sigmoid(out)\n",
    "\n",
    "            c_t = (forget * c) + (input * cell)\n",
    "            h_t = out * torch.tanh(c_t)\n",
    "\n",
    "            # Attention section\n",
    "            hidden_t, output = self.att(h_t, context, torch.eq(mask, 0))\n",
    "            hidden_t = torch.tanh(self.hidden_out(torch.cat((hidden_t, h_t), 1)))\n",
    "\n",
    "            return hidden_t, c_t, output\n",
    "\n",
    "        # Recurrence loop\n",
    "        for _ in range(input_length):\n",
    "            h_t, c_t, outs = step(decoder_input, hidden)\n",
    "            hidden = (h_t, c_t)\n",
    "\n",
    "            # Masking selected inputs\n",
    "            masked_outs = outs * mask\n",
    "\n",
    "            # Get maximum probabilities and indices\n",
    "            max_probs, indices = masked_outs.max(1)\n",
    "            one_hot_pointers = (runner == indices.unsqueeze(1).expand(-1, outs.size()[1])).float()\n",
    "\n",
    "            # Update mask to ignore seen indices\n",
    "            mask  = mask * (1 - one_hot_pointers)\n",
    "\n",
    "            # Get embedded inputs by max indices\n",
    "            embedding_mask = one_hot_pointers.unsqueeze(2).expand(-1, -1, self.embedding_dim).byte().bool()\n",
    "            decoder_input = embedded_inputs[embedding_mask.data].view(batch_size, self.embedding_dim)\n",
    "\n",
    "            outputs.append(outs.unsqueeze(0))\n",
    "            pointers.append(indices.unsqueeze(1))\n",
    "\n",
    "        outputs = torch.cat(outputs).permute(1, 0, 2)\n",
    "        pointers = torch.cat(pointers, 1)\n",
    "\n",
    "        return (outputs, pointers), hidden\n",
    "\n",
    "\n",
    "\n",
    "class PointerNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Pointer-Net\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embedding_dim,          \n",
    "                 hidden_dim,                   \n",
    "                 lstm_layers,                   \n",
    "                 dropout,                       \n",
    "                 bidir=False):                  \n",
    "        \"\"\"\n",
    "        Initiate Pointer-Net\n",
    "\n",
    "        :param int embedding_dim: Number of embbeding channels\n",
    "        :param int hidden_dim: Encoders hidden units\n",
    "        :param int lstm_layers: Number of layers for LSTMs\n",
    "        :param float dropout: Float between 0-1\n",
    "        :param bool bidir: Bidirectional\n",
    "        \"\"\"\n",
    "\n",
    "        super(PointerNet, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.bidir = bidir\n",
    "        self.embedding = nn.Linear(2, embedding_dim)\n",
    "        \n",
    "        self.encoder = Encoder(embedding_dim,\n",
    "                               hidden_dim,\n",
    "                               lstm_layers,\n",
    "                               dropout,\n",
    "                               bidir)\n",
    "        \n",
    "        self.decoder = Decoder(embedding_dim, hidden_dim)\n",
    "        self.decoder_input0 = Parameter(torch.FloatTensor(embedding_dim), requires_grad=False)\n",
    "\n",
    "        # Initialize decoder_input0\n",
    "        nn.init.uniform_(self.decoder_input0, -1, 1)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        \"\"\"\n",
    "        PointerNet - Forward-pass\n",
    "\n",
    "        :param Tensor inputs: Input sequence\n",
    "        :return: Pointers probabilities and indices\n",
    "        \"\"\"\n",
    "\n",
    "        batch_size = inputs.size(0)\n",
    "        input_length = inputs.size(1)\n",
    "\n",
    "        \"\"\"\n",
    "        请在此处将代码补充完整\n",
    "        \"\"\"\n",
    "\n",
    "        return  outputs, pointers\n",
    "\n",
    "train_size = 10000\n",
    "val_size = 1000\n",
    "test_size = 10000\n",
    "p_batch_size = 256\n",
    "\n",
    "nof_epoch = 50\n",
    "plr = 0.0001\n",
    "\n",
    "gpu = True\n",
    "\n",
    "nof_points = 5\n",
    "\n",
    "embedding_size = 128\n",
    "hiddens = 512\n",
    "nof_lstms = 2\n",
    "dropout = 0.0\n",
    "bidir = True\n",
    "\n",
    "if gpu and torch.cuda.is_available():\n",
    "    USE_CUDA = True\n",
    "    print('Using GPU, %i devices.' % torch.cuda.device_count())\n",
    "else:\n",
    "    USE_CUDA = False\n",
    "\n",
    "model = PointerNet(embedding_size,\n",
    "                   hiddens,\n",
    "                   nof_lstms,\n",
    "                   dropout,\n",
    "                   bidir)\n",
    "\n",
    "dataset = TSPDataset(train_size,\n",
    "                     nof_points)\n",
    "\n",
    "dataloader = DataLoader(dataset,\n",
    "                        batch_size=p_batch_size,\n",
    "                        shuffle=True,\n",
    "                        num_workers=4)\n",
    "\n",
    "if USE_CUDA:\n",
    "    model.cuda()\n",
    "    net = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "CCE = torch.nn.CrossEntropyLoss()\n",
    "model_optim = optim.Adam(filter(lambda p: p.requires_grad,\n",
    "                                    model.parameters()),\n",
    "                                    lr=plr)\n",
    "\n",
    "losses = []\n",
    "\n",
    "for epoch in range(nof_epoch):\n",
    "    batch_loss = []\n",
    "    iterator = tqdm(dataloader, unit='Batch')\n",
    "\n",
    "    for i_batch, sample_batched in enumerate(iterator):\n",
    "        iterator.set_description('Batch %i/%i' % (epoch + 1, nof_epoch))\n",
    "\n",
    "        train_batch = Variable(sample_batched['Points'])\n",
    "        target_batch = Variable(sample_batched['Solution'])\n",
    "\n",
    "        if USE_CUDA:\n",
    "            train_batch = train_batch.cuda()\n",
    "            target_batch = target_batch.cuda()\n",
    "\n",
    "        o, p = model(train_batch)\n",
    "        o = o.contiguous().view(-1, o.size()[-1])\n",
    "\n",
    "        target_batch = target_batch.view(-1)\n",
    "\n",
    "        loss = CCE(o, target_batch)\n",
    "\n",
    "        losses.append(loss.item())\n",
    "        batch_loss.append(loss.item())\n",
    "\n",
    "        model_optim.zero_grad()\n",
    "        loss.backward()\n",
    "        model_optim.step()\n",
    "\n",
    "        iterator.set_postfix(loss='{}'.format(loss.item()))\n",
    "\n",
    "    iterator.set_postfix(loss=np.average(batch_loss))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
